import argparse
import time

import gymnasium as gym
import numpy as np
import pygame


def parse_args():
    parser = argparse.ArgumentParser(description="Human play pendulum environment")
    parser.add_argument(
        "--episodes", type=int, default=3, help="Number of episodes to play"
    )
    parser.add_argument(
        "--max_steps", type=int, default=1000, help="Maximum steps per episode"
    )
    parser.add_argument(
        "--delay", type=float, default=0.02, help="Delay between frames in seconds"
    )
    parser.add_argument("--seed", type=int, default=42, help="Random seed")
    return parser.parse_args()


def main():
    args = parse_args()

    # 初始化pygame
    pygame.init()
    pygame.display.set_caption("Pendulum Human Control")

    # 创建环境
    env = gym.make("Pendulum-v1", render_mode="human", max_episode_steps=args.max_steps)
    action_bound = env.action_space.high[0]  # Pendulum的动作范围是[-2, 2]

    print("\n=== Pendulum Human Player ===")
    print("Controls:")
    print("  LEFT ARROW:  Apply torque to the left")
    print("  RIGHT ARROW: Apply torque to the right")
    print("  UP/DOWN:     Increase/decrease torque magnitude")
    print("  SPACE:       Reset to zero torque")
    print("  Q:           Quit the current episode")
    print("  ESC:         Exit the game")
    print("\nTry to swing the pendulum upright and balance it!")

    for episode in range(args.episodes):
        # 重置环境
        state = env.reset(seed=args.seed + episode)[0]
        total_reward = 0
        steps = 0
        done = False
        current_action = np.array([0.0])  # 初始为0扭矩
        action_magnitude = 1.0  # 初始扭矩大小为1

        print(f"\nEpisode {episode + 1}/{args.episodes} starting...")

        while not done and steps < args.max_steps:
            # 处理用户输入
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    env.close()
                    pygame.quit()
                    return
                elif event.type == pygame.KEYDOWN:
                    if event.key == pygame.K_ESCAPE:
                        env.close()
                        pygame.quit()
                        return
                    elif event.key == pygame.K_q:
                        done = True

            # 获取按键状态
            keys = pygame.key.get_pressed()

            # 更新扭矩大小
            if keys[pygame.K_UP]:
                action_magnitude = min(action_magnitude + 0.1, action_bound)
                print(f"Torque magnitude: {action_magnitude:.1f}")
            elif keys[pygame.K_DOWN]:
                action_magnitude = max(action_magnitude - 0.1, 0.1)
                print(f"Torque magnitude: {action_magnitude:.1f}")

            # 控制扭矩方向
            if keys[pygame.K_LEFT]:
                current_action = np.array([-action_magnitude])
            elif keys[pygame.K_RIGHT]:
                current_action = np.array([action_magnitude])
            elif keys[pygame.K_SPACE]:
                current_action = np.array([0.0])

            # 执行动作
            next_state, reward, terminated, truncated, _ = env.step(current_action)
            done = terminated or truncated

            total_reward += reward
            steps += 1
            state = next_state

            # 显示当前状态信息
            angle = np.arctan2(state[1], state[0])
            angle_deg = np.degrees(angle) % 360
            print(
                f"\rStep: {steps}, Angle: {angle_deg:.1f}°, Action: {current_action[0]:.2f}, "
                f"Reward: {reward:.2f}, Total Reward: {total_reward:.2f}",
                end="",
            )

            # 控制帧率
            time.sleep(args.delay)

        print(f"\nEpisode {episode + 1} finished with total reward: {total_reward:.2f}")

    env.close()
    pygame.quit()
    print("\nGame over. Thanks for playing!")


if __name__ == "__main__":
    main()
